import cv2
import time
import numpy as np
import torch
import torchvision
import torch.nn.functional as F
from rknnlite.api import RKNNLite



from coco_utils import COCO_test_helper


OBJ_THRESH = 0.25
NMS_THRESH = 0.45
MAX_DETECT = 300

IMG_SIZE = (640, 640)  # (width, height), such as (1280, 736)

CLASSES = ('Tactile Paving', 'Zebra Crossing')


class Colors:
    # Ultralytics color palette https://ultralytics.com/
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
                '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def filter_boxes(boxes, box_confidences, box_class_probs, seg_part):
    """Filter boxes with object threshold.
    """
    box_confidences = box_confidences.reshape(-1)
    candidate, class_num = box_class_probs.shape

    class_max_score = np.max(box_class_probs, axis=-1)
    classes = np.argmax(box_class_probs, axis=-1)

    _class_pos = np.where(class_max_score * box_confidences >= OBJ_THRESH)
    scores = (class_max_score * box_confidences)[_class_pos]

    boxes = boxes[_class_pos]
    classes = classes[_class_pos]
    seg_part = (seg_part * box_confidences.reshape(-1, 1))[_class_pos]

    return boxes, classes, scores, seg_part


def dfl(position):
    # Distribution Focal Loss (DFL)
    x = torch.tensor(position)
    n, c, h, w = x.shape
    p_num = 4
    mc = c // p_num
    y = x.reshape(n, p_num, mc, h, w)
    y = y.softmax(2)
    acc_metrix = torch.tensor(range(mc)).float().reshape(1, 1, mc, 1, 1)
    y = (y * acc_metrix).sum(2)
    return y.numpy()


def box_process(position):
    grid_h, grid_w = position.shape[2:4]
    col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
    col = col.reshape(1, 1, grid_h, grid_w)
    row = row.reshape(1, 1, grid_h, grid_w)
    grid = np.concatenate((col, row), axis=1)
    stride = np.array([IMG_SIZE[1] // grid_h, IMG_SIZE[0] // grid_w]).reshape(1, 2, 1, 1)

    position = dfl(position)
    box_xy = grid + 0.5 - position[:, 0:2, :, :]
    box_xy2 = grid + 0.5 + position[:, 2:4, :, :]
    xyxy = np.concatenate((box_xy * stride, box_xy2 * stride), axis=1)

    return xyxy


def post_process(input_data):
    # input_data[0], input_data[4], and input_data[8] are detection box information
    # input_data[1], input_data[5], and input_data[9] are category score information
    # input_data[2], input_data[6], and input_data[10] are confidence score information
    # input_data[3], input_data[7], and input_data[11] are segmentation information
    # input_data[12] is the proto information
    proto = input_data[-1]
    boxes, scores, classes_conf, seg_part = [], [], [], []
    defualt_branch = 3
    pair_per_branch = len(input_data) // defualt_branch

    for i in range(defualt_branch):
        boxes.append(box_process(input_data[pair_per_branch * i]))
        classes_conf.append(input_data[pair_per_branch * i + 1])
        scores.append(np.ones_like(input_data[pair_per_branch * i + 1][:, :1, :, :], dtype=np.float32))
        seg_part.append(input_data[pair_per_branch * i + 3])

    def sp_flatten(_in):
        ch = _in.shape[1]
        _in = _in.transpose(0, 2, 3, 1)
        return _in.reshape(-1, ch)

    boxes = [sp_flatten(_v) for _v in boxes]
    classes_conf = [sp_flatten(_v) for _v in classes_conf]
    scores = [sp_flatten(_v) for _v in scores]
    seg_part = [sp_flatten(_v) for _v in seg_part]

    boxes = np.concatenate(boxes)
    classes_conf = np.concatenate(classes_conf)
    scores = np.concatenate(scores)
    seg_part = np.concatenate(seg_part)

    # filter according to threshold
    boxes, classes, scores, seg_part = filter_boxes(boxes, scores, classes_conf, seg_part)

    zipped = zip(boxes, classes, scores, seg_part)
    sort_zipped = sorted(zipped, key=lambda x: (x[2]), reverse=True)
    result = zip(*sort_zipped)

    max_nms = 30000
    n = boxes.shape[0]  # number of boxes
    if not n:
        return None, None, None, None
    elif n > max_nms:  # excess boxes
        boxes, classes, scores, seg_part = [np.array(x[:max_nms]) for x in result]
    else:
        boxes, classes, scores, seg_part = [np.array(x) for x in result]

    # nms
    nboxes, nclasses, nscores, nseg_part = [], [], [], []
    agnostic = 0
    max_wh = 7680
    c = classes * (0 if agnostic else max_wh)
    ids = torchvision.ops.nms(
        torch.tensor(boxes, dtype=torch.float32) + torch.tensor(c, dtype=torch.float32).unsqueeze(-1),
        torch.tensor(scores, dtype=torch.float32), NMS_THRESH)
    real_keeps = ids.tolist()[:MAX_DETECT]
    nboxes.append(boxes[real_keeps])
    nclasses.append(classes[real_keeps])
    nscores.append(scores[real_keeps])
    nseg_part.append(seg_part[real_keeps])

    if not nclasses and not nscores:
        return None, None, None, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)
    seg_part = np.concatenate(nseg_part)

    ph, pw = proto.shape[-2:]
    proto = proto.reshape(seg_part.shape[-1], -1)
    seg_img = np.matmul(seg_part, proto)
    seg_img = sigmoid(seg_img)
    seg_img = seg_img.reshape(-1, ph, pw)

    seg_threadhold = 0.5

    # crop seg outside box
    seg_img = F.interpolate(torch.tensor(seg_img)[None], torch.Size([640, 640]), mode='bilinear', align_corners=False)[
        0]
    seg_img_t = _crop_mask(seg_img, torch.tensor(boxes))

    seg_img = seg_img_t.numpy()
    seg_img = seg_img > seg_threadhold
    return boxes, classes, scores, seg_img


def get_position(center_x, center_y, image_width, image_height):
    # Determine the direction of the object
    if center_y < image_height // 3:
        direction = "Left"
    elif center_y > 2 * image_height // 3:
        direction = "Right"
    else:
        direction = "Front"

    # Determine the distance of the object (based on the center_y position)
    if center_x < image_width // 3:
        distance = "Near"
    elif center_x > 2 * image_width // 3:
        distance = "Far"
    else:
        distance = "Around"

    return direction, distance


def draw(image, boxes, scores, classes):
    image_width, image_height = image.shape[:2]  # 获取图像的高和宽

    for box, score, cl in zip(boxes, scores, classes):
        top, left, right, bottom = [int(_b) for _b in box]

        # 计算底部中心点
        center_y = (top + right) // 2
        center_x = image_width - bottom
        print(center_x,center_y)

        # 获取目标的方向和距离
        direction, distance = get_position(center_x, center_y, image_width, image_height)

        # 输出目标类别及其位置
        print(f"{CLASSES[cl]} @ ({top} {left} {right} {bottom}) - {direction} {distance} %.3f" % (score))

        # 在图像中绘制目标框和标签
        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, f'{CLASSES[cl]} {score:.2f} {direction} {distance}',
                    (top, left - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)


def _crop_mask(masks, boxes):
    """
    "Crop" predicted masks by zeroing out everything not in the predicted bbox.
    Vectorized by Chong (thanks Chong).

    Args:
        - masks should be a size [h, w, n] tensor of masks
        - boxes should be a size [n, 4] tensor of bbox coords in relative point form
    """

    n, h, w = masks.shape
    x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(1,1,n)
    r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,w,1)
    c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]  # cols shape(h,1,1)

    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))


def merge_seg(image, seg_img, classes):
    color = Colors()
    for i in range(len(seg_img)):
        seg = seg_img[i]
        seg = seg.astype(np.uint8)
        seg = cv2.cvtColor(seg, cv2.COLOR_GRAY2BGR)
        seg = seg * color(classes[i])
        seg = seg.astype(np.uint8)
        image = cv2.add(image, seg)
    return image


def img_check(path):
    img_type = ['.jpg', '.jpeg', '.png', '.bmp']
    for _type in img_type:
        if path.endswith(_type) or path.endswith(_type.upper()):
            return True
    return False


if __name__ == '__main__':
    # 定义参数
    model_path = './model/best.rknn'  # 模型路径
    target = 'rk3588'  # RKNPU 目标平台
    IMG_SIZE = (640, 640)  # 模型输入尺寸，修改为实际尺寸

    # 初始化模型
    rknn_lite = RKNNLite()

    print('--> Load RKNN model')
    ret = rknn_lite.load_rknn(model_path)
    if ret != 0:
        print('Load RKNN model failed')
        exit(ret)
    print('done')

    print('--> Init runtime environment')
    if target == 'rk3588':
        ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0 | RKNNLite.NPU_CORE_1)
    else:
        ret = rknn_lite.init_runtime()

    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')

    # 初始化 COCO 测试助手
    co_helper = COCO_test_helper(enable_letter_box=True)

    # 打开摄像头
    cap = cv2.VideoCapture(0)  # 参数 0 表示使用默认摄像头
    if not cap.isOpened():
        print("Error: Could not open camera.")
        exit()

    # 计算帧率
    fps = 0
    frame_count = 0
    start_time = time.time()

    try:
        while True:
            # 从摄像头捕获画面
            ret, img_src = cap.read()
            if not ret:
                print("Error: Failed to capture frame.")
                break

            # 预处理图像
            pad_color = (0, 0, 0)
            img = co_helper.letter_box(im=img_src.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0]), pad_color=pad_color)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = np.expand_dims(img, 0)

            # 推理
            input_data = img
            outputs = rknn_lite.inference(inputs=[input_data])
            boxes, classes, scores, seg_img = post_process(outputs)

            # 初始化 img_p
            img_p = img_src.copy()  # 如果没有检测到目标，保留原始图像

            # 处理推理结果
            if boxes is not None:
                real_boxes = co_helper.get_real_box(boxes)
                real_segs = co_helper.get_real_seg(seg_img)
                img_p = merge_seg(img_src, real_segs, classes)  # 合并分割结果与原图

                # 绘制检测框
                draw(img_p, real_boxes, scores, classes)
            else:
                print("No objects detected.")

            # 计算帧率
            frame_count += 1
            elapsed_time = time.time() - start_time
            if elapsed_time > 1.0:  # 每秒更新一次帧率
                fps = frame_count / elapsed_time
                frame_count = 0
                start_time = time.time()

            # 显示帧率
            fps_text = f"FPS: {fps:.2f}"
            cv2.putText(img_p, fps_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # 显示结果
            cv2.imshow("Detection", img_p)

            # 按 'q' 键退出循环
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    finally:
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()

    rknn_lite.release()
    print('Processing completed.')



